// "Copyright [2021] <Copyright Shengkai Lin>"

#ifndef NETWORK_DEVICE_CPU_H_
#define NETWORK_DEVICE_CPU_H_

#include <vector>
#include <unordered_map>
#include "gpu.h"
#include "ai_sim.h"

namespace ai_sim {

struct TaskState {
    Task now_task;
    // The key of the map below is the flow id
    std::unordered_map<int, Flow> flows;
    std::vector<Flow> flows_queue;
};

class Gpu;

class Cpu : public omnetpp::cSimpleModule {
    friend class Gpu;
 private:
    int my_global_index;  // Used to decide if it is the master cpu
    int my_index;
    TaskState running_task;
    Gpu * my_gpu;

 protected:
    void initialize() override;
    void handleMessage(omnetpp::cMessage *msg) override;

    void alloc_flow(Flow);
    void handleNewTask(const TaskFlowVecMsg *);
    void handleReuqires(Flow&);
    void handleReceive(const FlowMsg *);
    void handleFinish(const FlowMsg *);
};

}  // namespace ai_sim


#endif  // NETWORK_DEVICE_CPU_H_
